"""中性可复用的动图生成工具

提供一个不含硬编码的函数 `generate_density_gif`，用于把若干年份的 GeoDataFrame 按列绘制为若干帧并合成 GIF。

设计要点：
- 使用 `src.plot.prepare_border_counties_population_data` 做数据准备（复用已有逻辑）。
- 参数化所有可配置项（列名、色带、dpi、figsize、colorbar bbox、字体名等）。
- 可选把每帧导出为 PNG（用于 QA）。
- 使用 Pillow 的 save_all 写入多帧 GIF，遇到异常回退到 imageio。
"""
from pathlib import Path
from typing import Iterable, Optional, Tuple
import io
import numpy as np
import pandas as pd
from PIL import Image as PILImage
import imageio
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.ticker import FuncFormatter

from src.plot import prepare_border_counties_population_data
from utils.plot_style import set_plot_style, set_chinese_font, force_chinese_font
import matplotlib.font_manager as fm
import json
from concurrent.futures import ProcessPoolExecutor, as_completed
import yaml
import multiprocessing


def _format_cb_ticks_identity(x, pos):
    """Format colorbar ticks using the original units (no division)."""
    # show integers when large, one decimal for small values
    try:
        if abs(x) >= 10:
            return f"{x:.0f}"
        return f"{x:.1f}"
    except Exception:
        return str(x)


def generate_density_gif(
    years: Iterable[int],
    out_path: str,
    population_parquet: str,
    column: str = 'mean',
    scope: str = 'county',
    cmap: str = 'Reds',
    dpi: int = 150,
    figsize: Tuple[int, int] = (10, 12),
    duration: float = 0.6,
    font_name: Optional[str] = None,
    save_frames: bool = False,
    frames_dir: str = 'figures/tmp_frames',
    verbose: bool = False,
    compress_percentile: int = 90,
    compress_gamma: float = 0.6,
) -> Tuple[str, int]:
    """生成动图并返回 (gif_path, frame_count)

    参数尽量中性化并可复用。
    """
    outp = Path(out_path)
    outp.parent.mkdir(parents=True, exist_ok=True)

    # 先收集每年的预处理结果并计算全局 vmin/vmax
    geos = {}
    vmin = None
    vmax = None
    for y in years:
        prep = prepare_border_counties_population_data(year=y, scope=scope, population_parquet=population_parquet)
        if not prep:
            if verbose:
                print(f"year {y}: no data, skip")
            continue
        geos[y] = prep
        geo = prep['geo']
        if column in geo.columns:
            ser = pd.to_numeric(geo[column], errors='coerce').fillna(0)
        elif 'mean' in geo.columns:
            ser = pd.to_numeric(geo['mean'], errors='coerce').fillna(0)
        else:
            ser = pd.Series(dtype=float)
        if not ser.empty:
            vals = ser.dropna().values
            if vals.size > 0:
                yr_min = float(vals.min())
                yr_max = float(vals.max())
                # 记录全局最小值
                if vmin is None or yr_min < vmin:
                    vmin = yr_min
                # 使用分位数上限以减少极端值对色标的影响（可配置百分位）
                p = float(np.percentile(vals, compress_percentile))
                if vmax is None or p > vmax:
                    vmax = p

    if not geos:
        raise RuntimeError('No data found for any year')

    if vmin is None or vmax is None:
        vmin, vmax = 0.0, 1.0
    if vmin == vmax:
        vmax = vmin + 1e-6

    frames = []
    frames_dir_path = Path(frames_dir)
    if save_frames:
        frames_dir_path.mkdir(parents=True, exist_ok=True)

    # 自动检测可用中文字体（当调用者未指定 font_name 时使用）
    preferred_cn = ['Noto Sans CJK SC', 'SimHei', 'PingFang', 'Heiti SC', 'WenQuanYi Zen Hei', 'Microsoft YaHei', 'STHeiti', 'SimSun']
    available_font_name = None
    try:
        names = {f.name for f in fm.fontManager.ttflist}
        for cand in preferred_cn:
            if cand in names:
                available_font_name = cand
                break
    except Exception:
        available_font_name = None

    for y in sorted(geos.keys()):
        # 我们将在后面并行渲染：将 GeoDataFrame 序列化为 __geo_interface__ 的 features
        prep = geos[y]
        geo = prep['geo']
        # 序列化要传输的轻量数据（features + columns）
        try:
            features = geo.__geo_interface__['features']
        except Exception:
            features = None
        prep['_serialized'] = {
            'features': features,
            'columns': {c: geo[c].fillna(None).tolist() for c in geo.columns if c != 'geometry'}
        }
        plot_cfg = prep.get('cfg_plot') or {}
        set_plot_style(plot_cfg)

        plot_figsize = prep.get('figsize', figsize)
        plot_dpi = prep.get('dpi', dpi)

        fig, ax = plt.subplots(figsize=figsize or plot_figsize, dpi=plot_dpi)
        # 选列
        if column in geo.columns:
            col = column
        else:
            col = 'mean' if 'mean' in geo.columns else '总人口'

        # 为缺失/无数据区域指定灰色填充，避免显示为透明或完全无色
        # 如果需要压缩色阶，使用 PowerNorm（gamma<1 会放大低值），否则使用线性归一化
        if compress_gamma and compress_gamma != 1.0:
            norm = mpl.colors.PowerNorm(gamma=compress_gamma, vmin=vmin, vmax=vmax)
        else:
            norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)

        try:
            geo.plot(column=col, cmap=cmap, linewidth=0.2, edgecolor='#444444', ax=ax, norm=norm, vmin=vmin, vmax=vmax, missing_kwds={"color": "#f7f7f7", "edgecolor": "#dddddd"})
        except Exception:
            # 兼容不同 geopandas 版本回退（某些版本不接受 norm）
            geo.plot(column=col, cmap=cmap, linewidth=0.2, edgecolor='#444444', ax=ax, vmin=vmin, vmax=vmax, missing_kwds={"color": "#f7f7f7", "edgecolor": "#dddddd"})
        ax.set_axis_off()

        mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
        mappable._A = []
        cax = fig.add_axes(prep['colorbar_bbox'])
        cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')
        cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_identity))
        cb.ax.tick_params(labelsize=9)
        # 统一图例标题
        cb.set_label('边境县乡镇人口密度', fontsize=11)

        # 字体处理：优先使用调用者传入的 font_name，其次使用 prepare 返回的建议字体，再尝试自动检测到的系统中文字体
        use_font = font_name or prep.get('font_family') or available_font_name
        if use_font:
            try:
                set_chinese_font(use_font)
                force_chinese_font(ax, use_font)
                for lab in cb.ax.get_xticklabels():
                    lab.set_fontfamily(use_font)
                cb.ax.xaxis.label.set_fontfamily(use_font)
            except Exception:
                pass

        # 年份注记（左上角）
        try:
            txt_kwargs = dict(color='white', fontsize=20, weight='bold', va='top', ha='left')
            ax.text(0.01, 0.99, str(y), transform=ax.transAxes, bbox=dict(facecolor='black', alpha=0.35, pad=4), **txt_kwargs)
        except Exception:
            pass

        buf = io.BytesIO()
        fig.savefig(buf, format='png', dpi=plot_dpi, bbox_inches='tight')
        plt.close(fig)
        # 并行渲染每帧：我们会把序列化数据发送到子进程进行绘图并返回 PNG bytes
        # 读取并行配置（优先使用仓库根目录下的 config.yaml）
        repo_root = Path(__file__).resolve().parents[1]
        cfg_path = repo_root / 'config.yaml'
        parallel = False
        n_threads = 1
        try:
            if cfg_path.exists():
                with open(cfg_path, 'r', encoding='utf-8') as fh:
                    cfg = yaml.safe_load(fh)
                parallel = bool(cfg.get('parallel', False))
                n_threads = int(cfg.get('n_threads', 1)) if cfg.get('n_threads') else (multiprocessing.cpu_count() or 1)
        except Exception:
            parallel = False
            n_threads = 1

        years_sorted = sorted(geos.keys())

        def _worker_render(args):
            # 在子进程中执行：接收 (year, prep_serialized, plot params)
            year, prep_ser, params = args
            import matplotlib as mpl
            mpl.use('Agg')
            import matplotlib.pyplot as plt
            import geopandas as gpd
            from shapely.geometry import shape
            import io, json

            # 重构 GeoDataFrame
        features = prep_ser.get('features')
        # properties -> DataFrame 保证与几何数量一致
        props = [feat.get('properties', {}) for feat in features] if features else []
        geoms = [shape(feat['geometry']) for feat in features] if features else []
        df = pd.DataFrame(props)
        gdf = gpd.GeoDataFrame(df, geometry=geoms, crs=params.get('crs', None))

            plot_figsize = params.get('figsize')
            plot_dpi = params.get('dpi')
            column = params.get('column')
            cmap = params.get('cmap')
            vmin = params.get('vmin')
            vmax = params.get('vmax')
            compress_gamma = params.get('compress_gamma')
            colorbar_bbox = params.get('colorbar_bbox')

            fig, ax = plt.subplots(figsize=plot_figsize, dpi=plot_dpi)
            if column in gdf.columns:
                col = column
            else:
                col = 'mean' if 'mean' in gdf.columns else '总人口'

            if compress_gamma and compress_gamma != 1.0:
                norm = mpl.colors.PowerNorm(gamma=compress_gamma, vmin=vmin, vmax=vmax)
            else:
                norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)

            try:
                gdf.plot(column=col, cmap=cmap, linewidth=0.2, edgecolor='#444444', ax=ax, norm=norm, vmin=vmin, vmax=vmax, missing_kwds={"color": "#f7f7f7", "edgecolor": "#dddddd"})
            except Exception:
                gdf.plot(column=col, cmap=cmap, linewidth=0.2, edgecolor='#444444', ax=ax, vmin=vmin, vmax=vmax, missing_kwds={"color": "#f7f7f7", "edgecolor": "#dddddd"})
            ax.set_axis_off()

            mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
            mappable._A = []
            cax = fig.add_axes(colorbar_bbox)
            cb = fig.colorbar(mappable, cax=cax, orientation='horizontal')
            cb.ax.xaxis.set_major_formatter(FuncFormatter(_format_cb_ticks_identity))
            cb.ax.tick_params(labelsize=9)
            cb.set_label('边境县乡镇人口密度', fontsize=11)

            # 年份注记
            try:
                txt_kwargs = dict(color='white', fontsize=20, weight='bold', va='top', ha='left')
                ax.text(0.01, 0.99, str(year), transform=ax.transAxes, bbox=dict(facecolor='black', alpha=0.35, pad=4), **txt_kwargs)
            except Exception:
                pass

            buf = io.BytesIO()
            fig.savefig(buf, format='png', dpi=plot_dpi, bbox_inches='tight')
            plt.close(fig)
            buf.seek(0)
            data = buf.getvalue()
            buf.close()
            return year, data

        # prepare args list
        args_list = []
        for y in years_sorted:
            prep = geos[y]
            ser = prep.get('_serialized')
            params = {
                'figsize': prep.get('figsize', figsize),
                'dpi': prep.get('dpi', dpi),
                'column': column,
                'cmap': cmap,
                'vmin': vmin,
                'vmax': vmax,
                'compress_gamma': compress_gamma,
                'colorbar_bbox': prep.get('colorbar_bbox'),
                'crs': prep.get('geo').crs if prep.get('geo') is not None else None,
            }
            args_list.append((y, ser, params))

        results = []
        if parallel and n_threads > 1:
            # macOS M1/M2/M4 建议使用 spawn 启动方式，但 ProcessPoolExecutor 会自动处理
            with ProcessPoolExecutor(max_workers=n_threads) as ex:
                futures = {ex.submit(_worker_render, arg): arg[0] for arg in args_list}
                for fut in as_completed(futures):
                    try:
                        yr, data = fut.result()
                        results.append((yr, data))
                    except Exception:
                        pass
        else:
            # 串行渲染
            for arg in args_list:
                yr, data = _worker_render(arg)
                results.append((yr, data))

        # 按年份排序并转换为 PIL frames
        results_sorted = sorted(results, key=lambda x: x[0])
        for yr, data in results_sorted:
            buf = io.BytesIO(data)
            pil_img = PILImage.open(buf).convert('RGBA')
            frames.append(pil_img.copy())
            buf.close()
